# FILE: generate_final_maps_and_c_tables.py
import numpy as np
import cv2

# =============================================================================
# Function to format and print C-style arrays for the lookup tables
# =============================================================================
def generate_c_array(array_data, array_name, rows, cols, data_type, macro_prefix):
    """
    Generates and prints a C-style 2D array declaration and initialization.
    
    Args:
        array_data (numpy.ndarray): The 2D numpy array containing the map data.
        array_name (str): The name of the array in the C code (e.g., "mapx").
        rows (int): Number of rows for the C array.
        cols (int): Number of columns for the C array.
        data_type (str): The data type for the C array (e.g., "float", "int").
        macro_prefix (str): The C macro to put before the declaration.
    """
    print(f"// Lookup table for {array_name}, dimensions: [{rows}][{cols}]")
    print(f"{macro_prefix}({data_type} {array_name}[{rows}][{cols}], 64) = {{")
    
    for i, row in enumerate(array_data):
        line = "{"
        line += ",".join([f"{val:.2f}" if data_type == "float" else str(int(val)) for val in row])
        line += "},"
        # Add a comment to easily identify the row number in the C code
        print(f"  {line} // Row {i}")
        
    print("};")
    print("\n" * 2) # Add some space before the next table

# =============================================================================
# PART 1 & 2: INPUT DATA and H_inv CALCULATION (No changes)
# =============================================================================
K = np.array([[98.5693, 0, 98.3482], [0, 98.7475, 56.5980], [0, 0, 1.0000]], dtype=np.float32)
rvec = np.array([[2.15770273], [0.00590957], [-0.0268242]], dtype=np.float32)
tvec = np.array([[-103.24288383], [129.7013359], [273.72416796]], dtype=np.float32)

# !!! IMPORTANT: Define the input image size of YOUR camera !!!
source_image_width = 188
source_image_height = 120

R, _ = cv2.Rodrigues(rvec)
H = K @ np.hstack((R[:, 0:2], tvec))
H_inv = np.linalg.inv(H)

# =============================================================================
# PART 3: AUTOMATIC CALCULATION FOR FINAL BIRD'S-EYE VIEW (No changes)
# =============================================================================
final_output_width_px = 376
final_output_height_px = 240
target_track_width_px = 46.0
real_track_width_mm = 450.0
pixels_per_mm = target_track_width_px / real_track_width_mm
track_center_px_current = (172.0 + 217.0) / 2.0
image_center_px = final_output_width_px / 2.0
pixel_offset = track_center_px_current - image_center_px
offset_x_mm = pixel_offset / pixels_per_mm
world_height_mm = final_output_height_px / pixels_per_mm
offset_y_mm = world_height_mm / 2.0
world_width_mm = final_output_width_px / pixels_per_mm
world_center_x = offset_x_mm
world_center_y = offset_y_mm
world_x_min = world_center_x - world_width_mm / 2.0
world_x_max = world_center_x + world_width_mm / 2.0
world_y_min = world_center_y - world_height_mm / 2.0
world_y_max = world_center_y + world_height_mm / 2.0
src_world_points = np.float32([
    [world_x_min, world_y_max], [world_x_max, world_y_max],
    [world_x_max, world_y_min], [world_x_min, world_y_min]
])
dst_bev_points = np.float32([
    [0, 0], [final_output_width_px - 1, 0],
    [final_output_width_px - 1, final_output_height_px - 1], [0, final_output_height_px - 1]
])
M_world_to_bev = cv2.getPerspectiveTransform(src_world_points, dst_bev_points)
final_transform_matrix = M_world_to_bev @ H_inv

# =============================================================================
# PART 4: GENERATE LOOKUP TABLES USING FORWARD MAPPING
# This new section generates the tables that your C code logic requires.
# =============================================================================

print("\n" + "="*60)
print("  Generating FORWARD MAPPING Lookup Tables for C Code")
print("="*60 + "\n")
print(f"Source Image Size: {source_image_width}x{source_image_height}")
print(f"Target Bird's-Eye View Size: {final_output_width_px}x{final_output_height_px}")
print(f"Output Table Dimensions: [{source_image_height}][{source_image_width}]")
print("\n" + "-"*60)


# 1. Create a grid of coordinates for every pixel in the SOURCE (188x120) image.
x_coords_src, y_coords_src = np.meshgrid(np.arange(source_image_width), np.arange(source_image_height))

# 2. Flatten the grid into a list of homogenous coordinates [x, y, 1].
src_pixels_flat = np.vstack((x_coords_src.flatten(), y_coords_src.flatten(), np.ones(x_coords_src.size)))

# 3. Apply the FORWARD transformation matrix (final_transform_matrix).
#    This calculates where each source pixel lands in the bird's-eye view.
bev_pixels_homogeneous = final_transform_matrix @ src_pixels_flat

# 4. Convert from homogeneous coordinates [u, v, w] back to 2D coordinates [u, v]
#    by dividing by the third coordinate (w).
bev_pixels_flat = bev_pixels_homogeneous[0:2, :] / bev_pixels_homogeneous[2, :]

# 5. Reshape the results back into two 120x188 maps.
#    mapx will contain the u-coordinates (X-axis) for the bird's-eye view.
#    mapy will contain the v-coordinates (Y-axis) for the bird's-eye view.
mapx_forward = bev_pixels_flat[0, :].reshape((source_image_height, source_image_width)).astype(np.float32)
mapy_forward = bev_pixels_flat[1, :].reshape((source_image_height, source_image_width)).astype(np.float32)


# 6. Generate and print the C code for the new mapx and mapy lookup tables.
print("Copy and paste the following code into your sj_camera_param.c file.")
print("IMPORTANT: Remember to change the array declarations in both .c and .h files to [120][188]!\n")

generate_c_array(mapx_forward, "mapx", source_image_height, source_image_width, "float", "AT_SDRAM_SECTION_ALIGN_INIT")
generate_c_array(mapy_forward, "mapy", source_image_height, source_image_width, "float", "AT_SDRAM_SECTION_ALIGN_INIT")

print("-" * 60)
print("Forward mapping table generation complete.")
print("="*60 + "\n")